package be.rivendale.renderer.mk2;

import be.rivendale.renderer.mk2.data.Rectangle;
import be.rivendale.renderer.mk2.data.Segment;

import java.util.*;

public class LoadBalancer implements Consumer {
//    private List<Producer> producers = new ArrayList<Producer>();
    private Map<Producer, ProducerStatistics> producers = new HashMap<Producer, ProducerStatistics>();

    @Override
    public void onConsumerStarted() {
    }

    @Override
    public void onSegmentComplete(Segment segment) {
        updateStatistics(segment.getRectangle());
    }

    private void updateStatistics(Rectangle rectangle) {
        int totalComplete = 0;
        for (Producer producer : producers.keySet()) {
            ProducerStatistics statistics = producers.get(producer);
            if(statistics.updatePendingWork(rectangle)) {
                producerComplete(producer, statistics);
            }
            if(statistics.isComplete()) {
                totalComplete++;
            }
        }
        if(totalComplete == producers.size()) {
            onAllStatisticsGathered();
        }
    }

    private void producerComplete(Producer producer, ProducerStatistics statistics) {
    }

    private void onAllStatisticsGathered() {
        double sumOfAverageTimesPerSegment = getSumOfAverageTimesPerSegment();
        for (Producer producer : producers.keySet()) {
            ProducerStatistics statistics = producers.get(producer);
            double ratio = statistics.getTotalTimePerNode() / sumOfAverageTimesPerSegment;
            statistics.setRatio(1 / ratio);
            System.out.format("\t - Producer '%s' completed it's %d segments in %.2fms at an average of %.2fms per node. This yields a unnormalized ratio of %.2f\n",
                    producer,
                    statistics.getInitialWorkSize(),
                    statistics.getTotalTime(),
                    statistics.getTotalTimePerNode(),
                    statistics.getRatio());
        }
        normalizeRatios();
    }

    private void normalizeRatios() {
        double sum = getSumOfRatios();
        for (ProducerStatistics producerStatistics : producers.values()) {
            producerStatistics.setRatio(producerStatistics.getRatio() / sum);
        }
    }

    private double getSumOfRatios() {
        double sum = 0;
        for (ProducerStatistics producerStatistics : producers.values()) {
            sum += producerStatistics.getRatio();
        }
        return sum;
    }

    private double getSumOfAverageTimesPerSegment() {
        double sumOfAverageTimesPerNode = 0;
        for (ProducerStatistics statistics : producers.values()) {
            sumOfAverageTimesPerNode += statistics.getTotalTimePerNode();
        }
        return sumOfAverageTimesPerNode;
    }

    public void dispatch(List<Rectangle> rectangles) {
        Collections.shuffle(rectangles);

        double fromIndex = 0;
        for (Map.Entry<Producer, ProducerStatistics> entry : producers.entrySet()) {
            double ratio = getPreviousRatio(entry);
            double howMany = rectangles.size() * ratio;
            double toIndex = fromIndex + howMany;
            ArrayList<Rectangle> rectanglesForThisProducer = new ArrayList<Rectangle>(rectangles.subList((int)Math.round(fromIndex), (int)Math.round(toIndex)));
            fromIndex = toIndex;
            entry.setValue(new ProducerStatistics(rectanglesForThisProducer));
            entry.getKey().offerWork(rectanglesForThisProducer);
        }
    }

    private double getPreviousRatio(Map.Entry<Producer, ProducerStatistics> entry) {
        if(entry.getValue() == null) {
            // First time, no statistics yet, code should be improved badly!!!
            return 1.0 / producers.size();
        } else {
            return entry.getValue().getRatio();
        }
    }

    public void addProducer(Producer producer) {
        producers.put(producer, null);
    }
}
